import pandas as pd
from langchain_ollama import OllamaLLM
import time

def extract_labels(text):
    # 初始化Ollama的Qwen模型
    llm = OllamaLLM(model="qwen2.5")
    
    prompt = f"""
    请从以下文本中提取3个最重要的标签词，用逗号分隔：
    {text}
    仅返回3个标签词，用逗号分隔，不要其他任何内容。
    """
    
    try:
        response = llm.invoke(prompt)
        labels = [label.strip() for label in response.split(',')]
        # 确保返回3个标签
        while len(labels) < 3:
            labels.append('')
        return labels[:3]
    except Exception as e:
        print(f"处理文本时出错: {e}")
        return ['', '', '']

def process_csv():
    # 读取CSV文件
    df = pd.read_csv('scene_data_20241110_120625.csv')
    total_rows = len(df)
    
    # 创建新的标签列
    df['label1'] = ''
    df['label2'] = ''
    df['label3'] = ''
    
    # 处理每一行的第一列文本
    for idx, row in df.iterrows():
        text = str(row.iloc[0])
        labels = extract_labels(text)
        
        # 更新标签列
        df.at[idx, 'label1'] = labels[0]
        df.at[idx, 'label2'] = labels[1]
        df.at[idx, 'label3'] = labels[2]
        
        # 显示进度
        print(f"处理进度: {idx+1}/{total_rows} ({((idx+1)/total_rows*100):.1f}%)")
        
        # 每处理5行保存一次，避免数据丢失
        if (idx + 1) % 5 == 0:
            df.to_csv('scene_data_processed.csv', index=False)
            print("已保存当前进度...")

    # 最终保存
    df.to_csv('scene_data_processed.csv', index=False)
    print("处理完成！结果已保存至 scene_data_processed.csv")

if __name__ == "__main__":
    process_csv()
